#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "jzp";

import numpy as np;
import pandas as pd;
from pandas import Series, DataFrame, DatetimeIndex;

'''
数据处理
'''

# numpy test
def np_test():
    print("---- 构建数组 ----");
    l = list(range(1,10));
    print(l);
    array = np.array(l);
    print(array);
    array = np.array([list(range(1, 10)), list(range(11,20))]);
    print(array);

    array = np.arange(20).reshape(4, 5);
    print(array);
    array = np.arange(1, 10);
    print(array);

    print("---- 广播 ----");
    print(l[0:3]);
    print(l[-3:]);
    array = np.arange(1, 10);
    print(array[0:7:2]);
    array = np.arange(20).reshape(4, 5);
    print(array);
    print(array[[0,1,2], [0, 1, 0]]); # (0, 0), (1, 1), (2,0)
    print(array[array > 4]);

    print("---- 广播 ----");
    array1 = np.arange(1, 10);
    array2 = np.arange(11, 20);
    print(array1);
    print(array2);
    print(array1 + array2);
    array1 = np.arange(20).reshape(4, 5);
    array2 = np.arange(20).reshape(4, 5);
    print(array1);
    print(array2);
    print(array1 + array2);

    print("---- 排序 ----");
    array = np.arange(1, 10);
    print(array);
    print(np.sort(array));
    print(-np.sort(-array));
    array = np.arange(20).reshape(4, 5);
    print(-np.sort(-array, axis=0));

    print("---- 统计 ----");
    array = np.arange(1, 10);
    print(np.amin(array));
    print(np.amax(array));
    print(np.average(array));
    print(np.median(array));
    array = np.arange(20).reshape(4, 5);
    print(array);
    print(np.amin(array, axis=0));
    print(np.amax(array, axis=1));
    print(np.average(array, axis=1));
    print(np.median(array, axis=1));

    print("---- 数组转list ----");
    array = np.arange(1, 10);
    print(array.tolist());
    array = np.arange(20).reshape(4, 5);
    print(array.tolist());

# Series 数据结构
def series_test():
    print("---- 构造 ----");
    data_list = list(range(1, 10));
    index_list = list("index_%d" % index for index in range(1, 10));
    s = Series(data=data_list);
    print(s);
    s = Series(data=data_list, index=index_list);
    print(s);

    data_dict = {"第一季度" : 32, "第二季度" : 43, "第三季度" : 245, "第四季度" : 89};
    index_list = ["武汉公司", "中国公司", "地球公司", "宇宙公司"];
    s = Series(data=data_dict);
    print(s);
    s = Series(data=data_dict, index=index_list);
    print(s);

    print("---- 新增 ----");
    data_list = list(range(1, 10));
    index_list = list("index_%d" % index for index in range(1, 10));
    s = Series(data=data_list, index=index_list);
    s["index_10"] = 3232;
    print(s);

    print("---- 修改 ----");
    s["index_9"] = 1111;
    print(s);

    print("---- 删除 ----");
    print(s.drop("index_10"));
    print(s.drop(["index_9", "index_10"]));

    print("---- 查询 ----");
    print(s.index);
    print(s.index.values.tolist());
    print(s.values.tolist());
    print(s.get("index_4"), s["index_5"]);
    print(s[[1, 4]]);
    print(s[1:4]);

    print("---- 转numpy ----");
    print(np.asarray(s));
    print(np.asarray(s).tolist());

# DataFrame 数据结构
def dataFrame_test():
    print("---- 构造 ----");
    data_dict = {"第一季度":[321, 432, 54, 65, 54], "第二季度":[34, 65, 43, 56, 54],
                 "第三季度":[34, 65, 98, 456, 54], "第四季度":[34, 54, 67, 89, 54]};
    df = DataFrame(data=data_dict);
    print(df);

    data_list = [[321, 432, 54, 65], [34, 65, 43, 56], [34, 65, 98, 456],
                 [34, 54, 67, 89], [45, 23, 56, 765]];
    index_list = ["武汉公司", "中国公司", "地球公司", "太阳系公司", "土星共同体"];
    column_list = ["第一季度", "第二季度", "第三季度", "第四季度"];
    df = DataFrame(data = data_list);
    print(df);
    df = DataFrame(data = data_list, index=index_list, columns=column_list);
    print(df);

    print("---- 新增 ----");
    df["年度"] = Series(data=[432, 54, 43, 765, 45], index=index_list);
    print(df);
    df["合格"] = (df["第四季度"] > 100);
    print(df);

    print("---- 删除 ----");
    print(df.drop("土星共同体", axis=0));
    print(df.drop(["土星共同体", "太阳系公司"], axis=0));
    print(df.drop("年度", axis=1));
    print(df.drop(["年度", "合格"], axis=1));

    print("---- 查询 ----");
    print(df["第二季度"]);
    print(df[["第二季度", "第四季度"]]);
    print(df.loc["武汉公司"]);
    print(df.loc[["武汉公司", "地球公司"]]);
    print(df.loc["武汉公司": "地球公司"]);
    print(df[0:2]);
    print(df.loc["中国公司", ["第二季度", "第三季度"]]);
    print(df.loc[["中国公司", "地球公司"], ["第二季度", "第三季度"]]);

    print("---- 遍历 ----");
    for index, row in df.iterrows():
        print(index, row);

    for item in df.itertuples():
        print(item.Index, item.第二季度);

    for x, y in zip(df["第一季度"], df["第三季度"]):
        print(x, y);

    print("---- 重置索引 ----");
    print(df.reindex(["武汉公司1", "中国公司", "地球公司", "太阳系公司", "土星共同体"], fill_value=11, axis=0));

    print("---- 排序 ----");
    print(df.sort_index(axis=1, ascending=True));
    print(df.sort_values(by="第一季度", axis=0, ascending=False));
    print(df.sort_values(by="武汉公司", axis=1, ascending=False));

    print("---- apply ----");
    print(df["第二季度"]);
    print(df["第二季度"].apply(lambda item : item + 2));
    df["第二季度"] = df["第二季度"].apply(lambda item : item + 2);
    print(df);

# 初始化测试数据
def init_series_dateFrame():
    data_list = list(range(1, 6));
    index_list = ["武汉公司", "中国公司", "地球公司", "太阳系公司", "土星共同体"];
    s = Series(data=data_list, index=index_list);

    data_list = [[321, 432, 54, 65], [34, 65, 43, 56], [34, 65, 98, 456],
                 [34, 54, 67, 89], [45, 23, 56, 765]];
    index_list = ["武汉公司", "中国公司", "地球公司", "太阳系公司", "土星共同体"];
    column_list = ["第一季度", "第二季度", "第三季度", "第四季度"];
    df = DataFrame(data=data_list, index=index_list, columns=column_list);

    return s, df;

# 算数运算
def arithmetic_test():
    s1, df1 = init_series_dateFrame();
    s2, df2 = init_series_dateFrame();
    # df2["年度"] = Series(data=[432, 54, 43, 765, 45],
    #                    index=["武汉公司", "中国公司", "地球公司", "太阳系公司", "土星共同体"]);

    print(s1 + s2);
    print(df1 + df2);
    print(df1.add(s1, axis=0));

# 数据统计
def data_analysis_test():
    s, df = init_series_dateFrame();

    print(s.describe());
    print(df.describe());
    print(s.count());
    print(s.sum());
    print(df.sum(axis=0));

# 输出文件
def df_csv_test():
    s, df = init_series_dateFrame();
    df.to_csv("D:/temp/df.csv", encoding="utf-8-sig");

    df2 = pd.read_csv("D:/temp/df.csv", encoding="utf-8-sig");
    print(df2);
    index_list = df2["Unnamed: 0"].tolist();
    print(index_list);
    df2 = df2.drop("Unnamed: 0", axis=1);
    column_list = df2.columns.values.tolist();
    print(column_list);
    print(DataFrame(data=df2.values, columns=column_list, index=index_list));

# 时间序列对象
def datetimeIndex_test():
    dti = pd.date_range(start="20230501", end="20230513");
    print(dti);
    print(dti.values);
    print(dti.tolist());
    print(dti.to_series());
    print(dti.to_frame());

if __name__ == '__main__':
    df_csv_test();